import openslide
import easyocr
import os
import glob
import time
import csv

# --- 1. 初始化 ---

# 指定包含 SVS 文件的目录
wsi_directory = "/mnt/data1/zzy/LNM/CRCpro-wsi/2025_10_ALL/"

# 指定输出 CSV 文件的路径
output_csv_file = "ocr_results.csv"

# 初始化 OCR 读取器（只需要执行一次）
print("正在加载 EasyOCR 模型...")
# 您可以根据需要更改语言，例如 ['ch_sim', 'en']
reader = easyocr.Reader(["en"], gpu=True)  # 如果有GPU，使用 gpu=True
print("模型加载完毕。")

# 临时的标签图像文件名
temp_label_image = "temp_label.png"

# --- 2. 查找所有 SVS 文件 ---
svs_files = glob.glob(os.path.join(wsi_directory, "*.svs"))

if not svs_files:
    print(f"在目录 {wsi_directory} 中未找到任何 .svs 文件。")
    exit()

print(f"找到了 {len(svs_files)} 个 .svs 文件，开始处理...")
print(f"结果将保存到: {output_csv_file}")
print("=============================================")

start_time = time.time()

# --- 3. 打开 CSV 文件并写入表头 ---
# 使用 'w' (写入) 模式, newline='' 是 csv 模块的推荐做法
with open(output_csv_file, "w", newline="", encoding="utf-8") as csvfile:
    # 定义 CSV 写入器
    csv_writer = csv.writer(csvfile)

    # 写入表头
    csv_writer.writerow(
        ["file_path", "text_rank", "text", "confidence", "y_coordinate", "full_bbox"]
    )

    # --- 4. 遍历和处理每个文件 ---
    for i, svs_file_path in enumerate(svs_files):
        file_name = os.path.basename(svs_file_path)
        print(f"\n[{i+1}/{len(svs_files)}] 正在处理: {file_name}")

        try:
            slide = openslide.OpenSlide(svs_file_path)

            if "label" in slide.associated_images:
                label_image = slide.associated_images["label"]
                label_image.save(temp_label_image)

                result = reader.readtext(temp_label_image)

                if not result:
                    print("  > OCR 未在该标签中识别到任何文本。")
                    # 写入 CSV 记录
                    csv_writer.writerow(
                        [svs_file_path, 0, "No text recognized", 0.0, -1, "N/A"]
                    )
                    slide.close()
                    continue

                # 按 Y 坐标排序
                sorted_result = sorted(result, key=lambda item: item[0][0][1])

                # 提取最顶端的两个文本
                top_two_texts = sorted_result[:2]

                print(
                    f"  > 成功识别到 {len(result)} 个文本，提取最顶端的 {len(top_two_texts)} 个:"
                )

                # --- 5. 将结果写入 CSV ---
                for rank, (bbox, text, prob) in enumerate(top_two_texts, 1):
                    y_coord = bbox[0][1]  # 左上角的 Y 坐标
                    print(
                        f"    - 文本: {text:<20} | 置信度: {prob:.2f} | Y坐标: {y_coord}"
                    )
                    # 写入 CSV
                    csv_writer.writerow(
                        [svs_file_path, rank, text, f"{prob:.4f}", y_coord, str(bbox)]
                    )

            else:
                print("  > 未找到 'label' 关联图像。")
                # 写入 CSV 记录
                csv_writer.writerow(
                    [svs_file_path, 0, "No 'label' image found", 0.0, -1, "N/A"]
                )

            slide.close()

        except Exception as e:
            print(f"  ! 处理 {file_name} 时发生错误: {e}")
            # 写入 CSV 记录
            csv_writer.writerow([svs_file_path, 0, f"Error: {e}", 0.0, -1, "N/A"])
            if "slide" in locals() and slide:
                slide.close()

# --- 6. 清理 ---
if os.path.exists(temp_label_image):
    os.remove(temp_label_image)

end_time = time.time()
print("\n=============================================")
print(f"所有文件处理完毕。总耗时: {end_time - start_time:.2f} 秒。")
print(f"结果已成功保存到 {output_csv_file}")
